from tasks import Task,Task_type
import random


# 用户节点类
class Node_U:
    def __init__(self,
                 nodeID,                # 用户 ID
                 task_type: int,        # 任务类型
                 computing_re=None,     # 最大计算需求
                 bandwidth_re=None):    # 最大带宽需求
        """
        :param task_type: 0-8：枚举任务类型,
                          9：代表手动配置
        :param nodeID:节点id
        """

        # user 属性信息
        self.nodeID = nodeID
        self.computing_re = computing_re
        self.bandwidth_re = bandwidth_re
        self.task_type = task_type
        assert (0 <= self.task_type < 11)

        # 用户生成的任务队列
        self.task_list = []

    def generate_task(self, taskID):
        if self.task_type != 9:
            members = list(Task_type.__members__.values())
            # print(members[self.task_type])
            computing_re, bandwidth_re = members[self.task_type].value
        else:
            computing_re, bandwidth_re = self.computing_re, self.bandwidth_re
        task = Task(computing_re, bandwidth_re, taskID, self.nodeID)  #

        # self.task_list.append(task)

        return task

##
# 工作节点类
class Node_W2:
    def __init__(self, computing_ca, bandwidth_ca, nodeID, strategy_c, strategy_t):
        # worker 属性信息
        self.nodeID = nodeID                    # 唯一标识 ID
        self.strategy_t = strategy_t
        self.strategy_c = strategy_c

        self.computing_ca = computing_ca        # 计算能力
        self.bandwidth_ca = bandwidth_ca        # 传输能力

        self.task_process_list = []             # 计算任务队列
        self.task_transmit_list = []            # 传输任务队列
        self.task_num = 0
        self.task_process_num = 0
        self.task_transmit_num = 0

    # 给节点分配相关的内容
    def allocate_task(self, type, task):
        assert type == 'T' or type == 'C'   # 传输任务为 T （transmission）& 计算任务为 C（computing） 且任务只能是这两种类型，目前只分配计算任务
        if type == 'T':
            self.task_transmit_list.append(task)  # insert to transmission tasks' queue
        elif type == 'C':
            self.task_process_list.append(task)   # insert to computing tasks' queue

    # worker 处理计算任务的函数
    def process_computing_task(self):
        if len(self.task_process_list) > 0:                         # 判断是否有计算任务，有则处理，无则输出无任务
            remind_computing = self.computing_ca
            del_list = []
            if self.strategy_c == 'FIFO':
                for index, task in enumerate(self.task_process_list):   # FIFO 先来先服务
                    task.process_time = task.process_time + 1
                    if remind_computing > task.computing_remind:
                        remind_computing = remind_computing - task.computing_remind
                        task.computing_remind = 0
                        self.task_process_num+=task.computing_re
                        del_list.append(index)
                        # print(f'处理完成计算任务，任务ID：{task.taskID}')
                        # print(f'添加到传输任务列表，任务ID：{task.taskID}')

                        self.task_transmit_list.append(task)
                    else:
                        task.computing_remind = task.computing_remind - remind_computing
                        remind_computing = 0
                self.task_process_list = [value for index, value in enumerate(self.task_process_list) if index not in del_list]  #更新
        else:
            # print(f'节点{self.nodeID}无计算任务')
            pass

    def process_transmit_task(self):
        task_done = []
        if len(self.task_transmit_list) > 0:                        # 判断是否有传输任务，有则处理，无则输出无任务
            remind_bandwidth = self.bandwidth_ca
            del_list = []
            if self.strategy_t == 'FIFO':
                for index, task in enumerate(self.task_transmit_list):  # FIFO 先来先服务
                    task.transmit_time += 1
                    if remind_bandwidth > task.bandwidth_remind:
                        remind_bandwidth -= task.bandwidth_remind
                        task.bandwidth_remind = 0
                        del_list.append(index)            # 处理完一个传输任务
                        # print(f'任务{task.taskID}处理完毕')
                        turnaround_time = task.transmit_time + task.process_time
                        re_bc = task.bandwidth_re+task.computing_remind
                        reward = re_bc/turnaround_time   #奖励1：带宽、算力单位处理速度作为奖励
                        # reward = -turnaround_time / re_bc  #奖励2，带权周转最小化
                        task_done.append((task.taskID,reward))
                        self.task_num +=1
                        self.task_transmit_num+=task.bandwidth_re
                    else:
                        task.bandwidth_remind -= remind_bandwidth
                        remind_bandwidth = 0

                self.task_transmit_list = [value for index, value in enumerate(self.task_transmit_list) if index not in del_list]  #更新



        else:
            # print(f'节点{self.nodeID}无传输任务')
            pass
        return task_done

    def get_node_state(self):
        computing_queue, bandwidth_queue = 0, 0
        for index, task in enumerate(self.task_process_list):
            computing_queue += task.computing_remind
        for index, task in enumerate(self.task_transmit_list):
            bandwidth_queue += task.bandwidth_remind

        return self.computing_ca, self.bandwidth_ca, computing_queue, bandwidth_queue
class Node_W1(Node_W2):
    def __init__(self,computing_ca, bandwidth_ca, nodeID, strategy_c, strategy_t,worker2:list):
        """

        :param computing_ca:
        :param bandwidth_ca:
        :param nodeID:
        :param strategy_c:
        :param strategy_t:
        :param worker2: 属于本节点的子计算节点对象列表
        """
        super().__init__(computing_ca, bandwidth_ca, nodeID, strategy_c, strategy_t)
        self.worker2 = worker2
        self.work_group = [self] + worker2

    def get_all_node_state(self):
        """
        计算本节点及子节点的计算能力和计算队列
        :return:
        """
        computing_queue, bandwidth_queue = 0, 0
        for index, task in enumerate(self.task_process_list):
            computing_queue += task.computing_remind
        for index, task in enumerate(self.task_transmit_list):
            bandwidth_queue += task.bandwidth_remind
        computing_ca, bandwidth_ca = self.computing_ca, self.bandwidth_ca
        for node in self.work_group:
            for index, task in enumerate(self.task_process_list):
                computing_queue += task.computing_remind
            for index, task in enumerate(self.task_transmit_list):
                bandwidth_queue += task.bandwidth_remind

            computing_ca += node.computing_ca
            bandwidth_ca += node.bandwidth_ca
        return computing_ca,bandwidth_ca,computing_queue,bandwidth_queue





